import json
import logging
from util.read_txt_file import ReadTxt
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
import os
from dotenv import load_dotenv
from catalog.dip_catalog import DipCatalog

logger = logging.getLogger(__name__)

class AgeSupplementaryCatalog:
    """年龄辅助目录处理类"""
    
    def __init__(self):
        self.age_knowledge = self._load_age_knowledge()
        
        # 初始化AI模型
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_BASE_URL")
        model_name = os.getenv("MODEL_NAME")
        
        self.llm = ChatOpenAI(
            temperature=0.1,
            max_tokens=300,
            base_url=base_url,
            model=model_name,
            api_key=api_key
        )
        
        # 简单的年龄分析提示
        self.age_prompt = ChatPromptTemplate.from_template(
            """你是年龄调整专家，请对患者进行年龄调整分析。

患者年龄：{patient_age}岁
患者姓名：{patient_name}
诊断编码：{diagnosis_code}
原始分组结果：{original_result}
年龄医疗知识库：{age_knowledge}
年龄辅助目录数据：{age_catalog_data}

请返回JSON格式的结果：

{{
  "分组名称": "[年龄辅助]加上年龄辅助目录中的病种名称",
  "分值": 调整后的数值,
  "调整说明": "调整原因和计算过程",
  "年龄调整完成": true/false 调整了才是true
}}

**重要要求**：
1. 必须返回有效的JSON格式
2. 分值必须是数值类型，不是字符串
3. 不要添加任何JSON之外的文字

返回JSON："""
        )
        self.json_parser = JsonOutputParser()
        self.age_chain = self.age_prompt | self.llm | self.json_parser

    #加载年龄医疗知识库
    def _load_age_knowledge(self):
        try:
            reader_txt = ReadTxt('年龄医疗知识库.txt')
            return reader_txt.load_txt_lines()
        except Exception as e:
            logger.error(f"加载年龄医疗知识库失败: {str(e)}")
            return ""

    #执行大模型分析年龄辅助目录
    def analyze_age_adjustment(self, patient_row, original_result,catalog_age):

        try:
            result = self.age_chain.invoke({
                "patient_age": str(patient_row['年龄']) + '岁',
                "patient_name": patient_row['患者姓名'],
                "diagnosis_code": patient_row['诊断编码'],
                "original_result": str(original_result),
                "age_catalog_data": catalog_age,
                "age_knowledge": self.age_knowledge
            })
            logger.info(f"年龄智能体分析结果: {result}")

            if isinstance(result, dict):
                return result
            elif isinstance(result, str):
                try:
                    return json.loads(result)
                except json.JSONDecodeError:
                    logger.error(f"床日智能体返回的不是有效JSON: {result}")
                    return {"错误": "JSON解析失败", "原始结果": result}
            else:
                return {"错误": "未知返回格式", "原始结果": str(result)}
        except Exception as e:
            logger.error(f"年龄智能体分析失败: {str(e)}")
            return {"错误": f"年龄调整失败: {str(e)}"}
    

